package org.hs.phd.odi.tracking.tracker.opencvmeanshift;

import java.awt.Polygon;
import java.awt.Rectangle;

import org.hs.phd.odi.tracking.tracker.Histogram;
import org.hs.phd.odi.tracking.tracker.TrackWindow;
import org.hs.phd.odi.tracking.tracker.Tracker;
import org.hs.phd.odi.tracking.tracker.histogram.Weighted3DHistogram;
import org.hs.phd.odi.tracking.tracker.kernel.EpanechnikovKernel;
import org.hs.phd.odi.tracking.tracker.trackwindow.RegularGridRectangularTrackWindow;

import com.googlecode.javacv.cpp.opencv_core;
import com.googlecode.javacv.cpp.opencv_core.CvBox2D;
import com.googlecode.javacv.cpp.opencv_core.CvPoint2D32f;
import com.googlecode.javacv.cpp.opencv_core.CvRect;
import com.googlecode.javacv.cpp.opencv_core.CvSize2D32f;
import com.googlecode.javacv.cpp.opencv_core.IplImage;
import com.googlecode.javacv.cpp.opencv_imgproc;
import com.googlecode.javacv.cpp.opencv_imgproc.CvConnectedComp;
import com.googlecode.javacv.cpp.opencv_video;

public class OpenCvMeanshiftTracker implements Tracker {

	private TrackWindow trackWindow;
	private boolean trackInitialized;
	private Histogram targetHistogram;

	
	public OpenCvMeanshiftTracker(){
		this( new RegularGridRectangularTrackWindow() );
	}
	
	protected OpenCvMeanshiftTracker(TrackWindow trackWindow){
		this.trackWindow = trackWindow;
		targetHistogram = new Weighted3DHistogram(32,32,16, new EpanechnikovKernel());
		trackInitialized = false;
	}
	

	protected opencv_core.CvBox2D findBestMatchTrack( IplImage backprojection, 
	  opencv_core.CvRect trackWindow )
	{
		
		CvConnectedComp comp = new CvConnectedComp();
		opencv_video.cvMeanShift(backprojection, trackWindow, 
				opencv_core.cvTermCriteria(opencv_core.CV_TERMCRIT_EPS | opencv_core.CV_TERMCRIT_ITER, 10, 1), comp);
		
		
		CvRect newRect = comp.rect();
		opencv_core.CvBox2D trackingResult = new opencv_core.CvBox2D();
		trackingResult.angle(0);
		trackingResult.center(new opencv_core.CvPoint2D32f(newRect.x() + newRect.width()/2.0,newRect.y() + newRect.height()/2.0));
		
		CvSize2D32f size = new opencv_core.CvSize2D32f();
		size.width(newRect.width());
		size.height(newRect.height());
		trackingResult.size(size);
		
		return trackingResult;

	}

	
	
	@Override
	public TrackWindow track(IplImage frame) {
		if(!trackInitialized) return null;

		IplImage backprojection = opencv_core.cvCreateImage(frame.cvSize(), frame.depth(), 1);
		targetHistogram.calculateBackProjection(frame, backprojection);		  
		//opencv_imgproc.cvThreshold(backprojection, backprojection, 50, 255, opencv_imgproc.CV_THRESH_TOZERO);
		
//		opencv_highgui.cvShowImage("backproj", backprojection);
//		opencv_highgui.cvWaitKey();
		 
		Rectangle trackWindowRect = trackWindow.getAsPolygon().getBounds();

		opencv_core.CvRect cvTrackWindow = new opencv_core.CvRect();
		cvTrackWindow.x(trackWindowRect.x);
		cvTrackWindow.y(trackWindowRect.y);
		cvTrackWindow.width(trackWindowRect.width);
		cvTrackWindow.height(trackWindowRect.height);
		
		CvBox2D bestMatchTrack = findBestMatchTrack( backprojection, cvTrackWindow);

		opencv_core.cvReleaseImage(backprojection);
		
		  
		trackWindow.moveToBestMatchingLocationFor( getAsPolygon(bestMatchTrack));
		return trackWindow;
	}

	private Polygon getAsPolygon(CvBox2D box) {
		Polygon result = new Polygon();
		
		opencv_core.CvPoint2D32f vertices = new opencv_core.CvPoint2D32f(4);
		opencv_imgproc.cvBoxPoints(box, vertices);

		for (int i = 0; i < 4; i++){
			CvPoint2D32f currentPoint = vertices.position(i);
			result.addPoint((int)currentPoint.x(), (int)currentPoint.y());
		}
		
		return result;
	}


	@Override
	public TrackWindow init(Polygon initialTrackShape, IplImage initialFrame) {
		trackInitialized = true;
		trackWindow.moveToBestMatchingLocationFor(initialTrackShape);
		targetHistogram.calculateFrom(initialFrame, trackWindow);
		return trackWindow;
	}

	@Override
	public TrackWindow getCurrentTrackWindow() {
		return trackWindow;
	}

	@Override
	public String getName() {
		return "Backprojection Meanshift Tracker";
	}

}
